{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0.0-alpha0\n",
      "sys.version_info(major=3, minor=5, micro=3, releaselevel='final', serial=0)\n",
      "matplotlib 3.0.3\n",
      "numpy 1.16.2\n",
      "pandas 0.24.2\n",
      "sklearn 0.21.2\n",
      "tensorflow 2.0.0-alpha0\n",
      "tensorflow.python.keras.api._v2.keras 2.2.4-tf\n"
     ]
    }
   ],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import sklearn\n",
    "import pandas as pd\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow import keras\n",
    "\n",
    "print(tf.__version__)\n",
    "print(sys.version_info)\n",
    "for module in mpl, np, pd, sklearn, tf, keras:\n",
    "    print(module.__name__, module.__version__)\n",
    "'''\n",
    "使用STML长短期记忆网络，做文本生成\n",
    "只要把原来的simplernn 改为stml\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1115394\n",
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You\n"
     ]
    }
   ],
   "source": [
    "# https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt\n",
    "input_filepath = \"./shakespeare.txt\"\n",
    "text = open(input_filepath, 'r').read()\n",
    "\n",
    "print(len(text))\n",
    "print(text[0:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "65\n",
      "['\\n', ' ', '!', '$', '&', \"'\", ',', '-', '.', '3', ':', ';', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']\n"
     ]
    }
   ],
   "source": [
    "# 1. generate vocab\n",
    "# 2. build mapping char->id\n",
    "# 3. data -> id_data\n",
    "# 4. abcd -> bcd<eos>\n",
    "\n",
    "vocab = sorted(set(text))\n",
    "print(len(vocab))\n",
    "print(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'f': 44, 'g': 45, 'W': 35, 'X': 36, 'G': 19, 'h': 46, 'N': 26, 'z': 64, 'j': 48, 'm': 51, 'e': 43, '\\n': 0, '!': 2, 'd': 42, 'U': 33, 't': 58, ':': 10, 'T': 32, 'S': 31, 'n': 52, 'Z': 38, '.': 8, 'k': 49, ' ': 1, 'P': 28, 'K': 23, ',': 6, 'F': 18, '&': 4, 'a': 39, 'R': 30, 'H': 20, 'w': 61, 'B': 14, 'u': 59, 'Y': 37, '3': 9, '$': 3, 'p': 54, '-': 7, 'A': 13, 's': 57, 'l': 50, \"'\": 5, 'O': 27, 'J': 22, 'L': 24, 'c': 41, 'Q': 29, 'V': 34, 'D': 16, 'v': 60, 'x': 62, ';': 11, 'E': 17, 'r': 56, 'q': 55, 'i': 47, 'b': 40, '?': 12, 'C': 15, 'o': 53, 'M': 25, 'I': 21, 'y': 63}\n"
     ]
    }
   ],
   "source": [
    "char2idx = {char:idx for idx, char in enumerate(vocab)}\n",
    "print(char2idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['\\n' ' ' '!' '$' '&' \"'\" ',' '-' '.' '3' ':' ';' '?' 'A' 'B' 'C' 'D' 'E'\n",
      " 'F' 'G' 'H' 'I' 'J' 'K' 'L' 'M' 'N' 'O' 'P' 'Q' 'R' 'S' 'T' 'U' 'V' 'W'\n",
      " 'X' 'Y' 'Z' 'a' 'b' 'c' 'd' 'e' 'f' 'g' 'h' 'i' 'j' 'k' 'l' 'm' 'n' 'o'\n",
      " 'p' 'q' 'r' 's' 't' 'u' 'v' 'w' 'x' 'y' 'z']\n"
     ]
    }
   ],
   "source": [
    "idx2char = np.array(vocab)\n",
    "print(idx2char)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[18 47 56 57 58  1 15 47 58 47]\n",
      "First Citi\n"
     ]
    }
   ],
   "source": [
    "text_as_int = np.array([char2idx[c] for c in text])\n",
    "print(text_as_int[0:10])\n",
    "print(text[0:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(18, shape=(), dtype=int64) F\n",
      "tf.Tensor(47, shape=(), dtype=int64) i\n",
      "tf.Tensor(\n",
      "[18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 14 43 44 53 56 43  1 61 43\n",
      "  1 54 56 53 41 43 43 42  1 39 52 63  1 44 59 56 58 46 43 56  6  1 46 43\n",
      " 39 56  1 51 43  1 57 54 43 39 49  8  0  0 13 50 50 10  0 31 54 43 39 49\n",
      "  6  1 57 54 43 39 49  8  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10\n",
      "  0 37 53 59  1], shape=(101,), dtype=int64)\n",
      "'First Citizen:\\nBefore we proceed any further, hear me speak.\\n\\nAll:\\nSpeak, speak.\\n\\nFirst Citizen:\\nYou '\n",
      "tf.Tensor(\n",
      "[39 56 43  1 39 50 50  1 56 43 57 53 50 60 43 42  1 56 39 58 46 43 56  1\n",
      " 58 53  1 42 47 43  1 58 46 39 52  1 58 53  1 44 39 51 47 57 46 12  0  0\n",
      " 13 50 50 10  0 30 43 57 53 50 60 43 42  8  1 56 43 57 53 50 60 43 42  8\n",
      "  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 18 47 56 57 58  6  1\n",
      " 63 53 59  1 49], shape=(101,), dtype=int64)\n",
      "'are all resolved rather to die than to famish?\\n\\nAll:\\nResolved. resolved.\\n\\nFirst Citizen:\\nFirst, you k'\n"
     ]
    }
   ],
   "source": [
    "def split_input_target(id_text):\n",
    "    \"\"\"\n",
    "    abcde -> abcd, bcde\n",
    "    \"\"\"\n",
    "    return id_text[0:-1], id_text[1:]\n",
    "\n",
    "char_dataset = tf.data.Dataset.from_tensor_slices(text_as_int)\n",
    "seq_length = 100\n",
    "seq_dataset = char_dataset.batch(seq_length + 1,\n",
    "                                 drop_remainder = True)\n",
    "for ch_id in char_dataset.take(2):\n",
    "    print(ch_id, idx2char[ch_id.numpy()])\n",
    "\n",
    "for seq_id in seq_dataset.take(2):\n",
    "    print(seq_id)\n",
    "    print(repr(''.join(idx2char[seq_id.numpy()])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 14 43 44 53 56 43  1 61 43\n",
      "  1 54 56 53 41 43 43 42  1 39 52 63  1 44 59 56 58 46 43 56  6  1 46 43\n",
      " 39 56  1 51 43  1 57 54 43 39 49  8  0  0 13 50 50 10  0 31 54 43 39 49\n",
      "  6  1 57 54 43 39 49  8  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10\n",
      "  0 37 53 59]\n",
      "[47 56 57 58  1 15 47 58 47 64 43 52 10  0 14 43 44 53 56 43  1 61 43  1\n",
      " 54 56 53 41 43 43 42  1 39 52 63  1 44 59 56 58 46 43 56  6  1 46 43 39\n",
      " 56  1 51 43  1 57 54 43 39 49  8  0  0 13 50 50 10  0 31 54 43 39 49  6\n",
      "  1 57 54 43 39 49  8  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10  0\n",
      " 37 53 59  1]\n",
      "[39 56 43  1 39 50 50  1 56 43 57 53 50 60 43 42  1 56 39 58 46 43 56  1\n",
      " 58 53  1 42 47 43  1 58 46 39 52  1 58 53  1 44 39 51 47 57 46 12  0  0\n",
      " 13 50 50 10  0 30 43 57 53 50 60 43 42  8  1 56 43 57 53 50 60 43 42  8\n",
      "  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 18 47 56 57 58  6  1\n",
      " 63 53 59  1]\n",
      "[56 43  1 39 50 50  1 56 43 57 53 50 60 43 42  1 56 39 58 46 43 56  1 58\n",
      " 53  1 42 47 43  1 58 46 39 52  1 58 53  1 44 39 51 47 57 46 12  0  0 13\n",
      " 50 50 10  0 30 43 57 53 50 60 43 42  8  1 56 43 57 53 50 60 43 42  8  0\n",
      "  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 18 47 56 57 58  6  1 63\n",
      " 53 59  1 49]\n"
     ]
    }
   ],
   "source": [
    "seq_dataset = seq_dataset.map(split_input_target)\n",
    "\n",
    "for item_input, item_output in seq_dataset.take(2):\n",
    "    print(item_input.numpy())\n",
    "    print(item_output.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "buffer_size = 10000\n",
    "\n",
    "seq_dataset = seq_dataset.shuffle(buffer_size).batch(\n",
    "    batch_size, drop_remainder=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0609 11:39:08.665987 140305539143424 tf_logging.py:161] <tensorflow.python.keras.layers.recurrent.UnifiedLSTM object at 0x7f9ac77b77f0>: Note that this layer is not optimized for performance. Please use tf.keras.layers.CuDNNLSTM for better performance on GPU.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding (Embedding)        (64, None, 256)           16640     \n",
      "_________________________________________________________________\n",
      "unified_lstm (UnifiedLSTM)   (64, None, 1024)          5246976   \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (64, None, 65)            66625     \n",
      "=================================================================\n",
      "Total params: 5,330,241\n",
      "Trainable params: 5,330,241\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "vocab_size = len(vocab)\n",
    "embedding_dim = 256\n",
    "rnn_units = 1024\n",
    "\n",
    "def build_model(vocab_size, embedding_dim, rnn_units, batch_size):\n",
    "    model = keras.models.Sequential([\n",
    "        keras.layers.Embedding(vocab_size, embedding_dim,\n",
    "                               batch_input_shape = [batch_size, None]),\n",
    "        keras.layers.LSTM(units = rnn_units,\n",
    "                          stateful = True,\n",
    "                          recurrent_initializer = 'glorot_uniform',\n",
    "                          return_sequences = True),\n",
    "        keras.layers.Dense(vocab_size),\n",
    "    ])\n",
    "    return model\n",
    "\n",
    "model = build_model(\n",
    "    vocab_size = vocab_size,\n",
    "    embedding_dim = embedding_dim,\n",
    "    rnn_units = rnn_units,\n",
    "    batch_size = batch_size)\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 100, 65)\n"
     ]
    }
   ],
   "source": [
    "for input_example_batch, target_example_batch in seq_dataset.take(1):\n",
    "    example_batch_predictions = model(input_example_batch)\n",
    "    print(example_batch_predictions.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[34]\n",
      " [62]\n",
      " [15]\n",
      " [29]\n",
      " [58]\n",
      " [52]\n",
      " [ 4]\n",
      " [33]\n",
      " [37]\n",
      " [ 9]\n",
      " [27]\n",
      " [29]\n",
      " [34]\n",
      " [64]\n",
      " [ 0]\n",
      " [20]\n",
      " [47]\n",
      " [58]\n",
      " [ 6]\n",
      " [60]\n",
      " [26]\n",
      " [ 3]\n",
      " [17]\n",
      " [48]\n",
      " [30]\n",
      " [58]\n",
      " [10]\n",
      " [62]\n",
      " [41]\n",
      " [ 2]\n",
      " [32]\n",
      " [47]\n",
      " [22]\n",
      " [26]\n",
      " [ 3]\n",
      " [57]\n",
      " [58]\n",
      " [36]\n",
      " [29]\n",
      " [11]\n",
      " [ 1]\n",
      " [59]\n",
      " [54]\n",
      " [16]\n",
      " [ 3]\n",
      " [18]\n",
      " [31]\n",
      " [ 6]\n",
      " [ 1]\n",
      " [44]\n",
      " [24]\n",
      " [38]\n",
      " [ 8]\n",
      " [21]\n",
      " [47]\n",
      " [27]\n",
      " [17]\n",
      " [26]\n",
      " [61]\n",
      " [13]\n",
      " [37]\n",
      " [61]\n",
      " [59]\n",
      " [54]\n",
      " [33]\n",
      " [25]\n",
      " [37]\n",
      " [46]\n",
      " [28]\n",
      " [10]\n",
      " [ 6]\n",
      " [22]\n",
      " [ 3]\n",
      " [49]\n",
      " [ 1]\n",
      " [28]\n",
      " [53]\n",
      " [50]\n",
      " [29]\n",
      " [21]\n",
      " [33]\n",
      " [58]\n",
      " [14]\n",
      " [54]\n",
      " [28]\n",
      " [22]\n",
      " [ 2]\n",
      " [61]\n",
      " [55]\n",
      " [43]\n",
      " [14]\n",
      " [36]\n",
      " [ 2]\n",
      " [18]\n",
      " [29]\n",
      " [27]\n",
      " [50]\n",
      " [43]\n",
      " [56]\n",
      " [40]], shape=(100, 1), dtype=int64)\n",
      "tf.Tensor(\n",
      "[34 62 15 29 58 52  4 33 37  9 27 29 34 64  0 20 47 58  6 60 26  3 17 48\n",
      " 30 58 10 62 41  2 32 47 22 26  3 57 58 36 29 11  1 59 54 16  3 18 31  6\n",
      "  1 44 24 38  8 21 47 27 17 26 61 13 37 61 59 54 33 25 37 46 28 10  6 22\n",
      "  3 49  1 28 53 50 29 21 33 58 14 54 28 22  2 61 55 43 14 36  2 18 29 27\n",
      " 50 43 56 40], shape=(100,), dtype=int64)\n"
     ]
    }
   ],
   "source": [
    "# random sampling.\n",
    "# greedy, random.\n",
    "sample_indices = tf.random.categorical(\n",
    "    logits = example_batch_predictions[0], num_samples = 1)\n",
    "print(sample_indices)\n",
    "# (100, 65) -> (100, 1)\n",
    "sample_indices = tf.squeeze(sample_indices, axis = -1)\n",
    "print(sample_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input:  \"omes not Rome, nor has Coriolanus\\nDeserved this so dishonour'd rub, laid falsely\\nI' the plain way of\"\n",
      "\n",
      "Output:  \"mes not Rome, nor has Coriolanus\\nDeserved this so dishonour'd rub, laid falsely\\nI' the plain way of \"\n",
      "\n",
      "Predictions:  'VxCQtn&UY3OQVz\\nHit,vN$EjRt:xc!TiJN$stXQ; upD$FS, fLZ.IiOENwAYwupUMYhP:,J$k PolQIUtBpPJ!wqeBX!FQOlerb'\n"
     ]
    }
   ],
   "source": [
    "print(\"Input: \", repr(\"\".join(idx2char[input_example_batch[0]])))\n",
    "print()\n",
    "print(\"Output: \", repr(\"\".join(idx2char[target_example_batch[0]])))\n",
    "print()\n",
    "print(\"Predictions: \", repr(\"\".join(idx2char[sample_indices])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 100)\n",
      "4.174524\n"
     ]
    }
   ],
   "source": [
    "def loss(labels, logits):\n",
    "    return keras.losses.sparse_categorical_crossentropy(\n",
    "        labels, logits, from_logits=True)\n",
    "\n",
    "model.compile(optimizer = 'adam', loss = loss)\n",
    "example_loss = loss(target_example_batch, example_batch_predictions)\n",
    "print(example_loss.shape)\n",
    "print(example_loss.numpy().mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/100\n",
      "172/172 [==============================] - 8s 44ms/step - loss: 2.6133\n",
      "Epoch 2/100\n",
      "172/172 [==============================] - 7s 39ms/step - loss: 1.9266\n",
      "Epoch 3/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 1.6686\n",
      "Epoch 4/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 1.5240\n",
      "Epoch 5/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 1.4352\n",
      "Epoch 6/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 1.3713\n",
      "Epoch 7/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 1.3192\n",
      "Epoch 8/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 1.2734\n",
      "Epoch 9/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 1.2309\n",
      "Epoch 10/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 1.1907\n",
      "Epoch 11/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 1.1530\n",
      "Epoch 12/100\n",
      "172/172 [==============================] - 6s 38ms/step - loss: 1.1166\n",
      "Epoch 13/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 1.0780\n",
      "Epoch 14/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 1.0417\n",
      "Epoch 15/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 1.0066\n",
      "Epoch 16/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.9715\n",
      "Epoch 17/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.9388\n",
      "Epoch 18/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.9093\n",
      "Epoch 19/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.8818\n",
      "Epoch 20/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.8509\n",
      "Epoch 21/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.8220\n",
      "Epoch 22/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.7966\n",
      "Epoch 23/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.7741\n",
      "Epoch 24/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.7498\n",
      "Epoch 25/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.7268\n",
      "Epoch 26/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.7022\n",
      "Epoch 27/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.6790\n",
      "Epoch 28/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.6588\n",
      "Epoch 29/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.6401\n",
      "Epoch 30/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.6254\n",
      "Epoch 31/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.6102\n",
      "Epoch 32/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.5924\n",
      "Epoch 33/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.5790\n",
      "Epoch 34/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.5654\n",
      "Epoch 35/100\n",
      "172/172 [==============================] - 6s 38ms/step - loss: 0.5510\n",
      "Epoch 36/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 0.5391\n",
      "Epoch 37/100\n",
      "172/172 [==============================] - 6s 38ms/step - loss: 0.5240\n",
      "Epoch 38/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 0.5122\n",
      "Epoch 39/100\n",
      "172/172 [==============================] - 7s 39ms/step - loss: 0.5005\n",
      "Epoch 40/100\n",
      "172/172 [==============================] - 6s 38ms/step - loss: 0.4902\n",
      "Epoch 41/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.4770\n",
      "Epoch 42/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.4691\n",
      "Epoch 43/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.4605\n",
      "Epoch 44/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 0.4523\n",
      "Epoch 45/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 0.4439\n",
      "Epoch 46/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 0.4357\n",
      "Epoch 47/100\n",
      "172/172 [==============================] - 6s 38ms/step - loss: 0.4300\n",
      "Epoch 48/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.4229\n",
      "Epoch 49/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 0.4146\n",
      "Epoch 50/100\n",
      "172/172 [==============================] - 7s 41ms/step - loss: 0.4110\n",
      "Epoch 51/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.4020\n",
      "Epoch 52/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3954\n",
      "Epoch 53/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.3885\n",
      "Epoch 54/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.3845\n",
      "Epoch 55/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3816\n",
      "Epoch 56/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3726\n",
      "Epoch 57/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.3686\n",
      "Epoch 58/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.3602\n",
      "Epoch 59/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3531\n",
      "Epoch 60/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.3448\n",
      "Epoch 61/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3394\n",
      "Epoch 62/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.3375\n",
      "Epoch 63/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3351\n",
      "Epoch 64/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3361\n",
      "Epoch 65/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3293\n",
      "Epoch 66/100\n",
      "172/172 [==============================] - 7s 39ms/step - loss: 0.3227\n",
      "Epoch 67/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3230\n",
      "Epoch 68/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.3207\n",
      "Epoch 69/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3188\n",
      "Epoch 70/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3161\n",
      "Epoch 71/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3134\n",
      "Epoch 72/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3123\n",
      "Epoch 73/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3097\n",
      "Epoch 74/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.3117\n",
      "Epoch 75/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3122\n",
      "Epoch 76/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3054\n",
      "Epoch 77/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.3018\n",
      "Epoch 78/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2964\n",
      "Epoch 79/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2974\n",
      "Epoch 80/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2939\n",
      "Epoch 81/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2902\n",
      "Epoch 82/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2851\n",
      "Epoch 83/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2800\n",
      "Epoch 84/100\n",
      "172/172 [==============================] - 7s 38ms/step - loss: 0.2774\n",
      "Epoch 85/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2761\n",
      "Epoch 86/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2755\n",
      "Epoch 87/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2790\n",
      "Epoch 88/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2757\n",
      "Epoch 89/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2740\n",
      "Epoch 90/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2738\n",
      "Epoch 91/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2746\n",
      "Epoch 92/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2731\n",
      "Epoch 93/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2731\n",
      "Epoch 94/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2745\n",
      "Epoch 95/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2699\n",
      "Epoch 96/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2665\n",
      "Epoch 97/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2648\n",
      "Epoch 98/100\n",
      "172/172 [==============================] - 6s 37ms/step - loss: 0.2648\n",
      "Epoch 99/100\n",
      "172/172 [==============================] - 7s 41ms/step - loss: 0.2663\n",
      "Epoch 100/100\n",
      "172/172 [==============================] - 6s 36ms/step - loss: 0.2628\n"
     ]
    }
   ],
   "source": [
    "output_dir = \"./text_generation_lstm3_checkpoints\"\n",
    "if not os.path.exists(output_dir):\n",
    "    os.mkdir(output_dir)\n",
    "checkpoint_prefix = os.path.join(output_dir, 'ckpt_{epoch}')\n",
    "checkpoint_callback = keras.callbacks.ModelCheckpoint(\n",
    "    filepath = checkpoint_prefix,\n",
    "    save_weights_only = True)\n",
    "\n",
    "epochs = 100\n",
    "history = model.fit(seq_dataset, epochs = epochs,\n",
    "                    callbacks = [checkpoint_callback])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'./text_generation_lstm3_checkpoints/ckpt_100'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.train.latest_checkpoint(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0609 11:49:51.541063 140305539143424 tf_logging.py:161] <tensorflow.python.keras.layers.recurrent.UnifiedLSTM object at 0x7f9ab047a550>: Note that this layer is not optimized for performance. Please use tf.keras.layers.CuDNNLSTM for better performance on GPU.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_1\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding_1 (Embedding)      (1, None, 256)            16640     \n",
      "_________________________________________________________________\n",
      "unified_lstm_1 (UnifiedLSTM) (1, None, 1024)           5246976   \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (1, None, 65)             66625     \n",
      "=================================================================\n",
      "Total params: 5,330,241\n",
      "Trainable params: 5,330,241\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model2 = build_model(vocab_size,\n",
    "                     embedding_dim,\n",
    "                     rnn_units,\n",
    "                     batch_size = 1)\n",
    "model2.load_weights(tf.train.latest_checkpoint(output_dir))\n",
    "model2.build(tf.TensorShape([1, None]))\n",
    "# start ch sequence A,\n",
    "# A -> model -> b\n",
    "# A.append(b) -> B\n",
    "# B(Ab) -> model -> c\n",
    "# B.append(c) -> C\n",
    "# C(Abc) -> model -> ...\n",
    "model2.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "All: anch's word: yet Juchio,\n",
      "For you'll proceed was wounds bid God call.\n",
      "\n",
      "CORIOLANUS:\n",
      "Know, old Juliet,\n",
      "Was my willing unch losk in dissolute a\n",
      "crown in Geople.\n",
      "\n",
      "Boy:\n",
      "Gragorr! what an you br\n",
      "The good man of world, yet one that groans?\n",
      "\n",
      "EDWARD:\n",
      "The king about he,\n",
      "Have made Padua, there is shr\n",
      "Deas cannot oft, you be it soil'd;\n",
      "Which, but I'll dry your judgment,\n",
      "Because that Romeo, ear the slim! this is thy hamb,\n",
      "More birth, singul, rule, in one of the Capelbur.\n",
      "\n",
      "Herd Servingman:\n",
      "By both give hunch, tear nighting ' nt as I.\n",
      "Ah, Juliet, quStagion ned Claudio,\n",
      "Show Edward's queen inchroptish, ESTAUnk'd one that so? What corrads over to readonWy with master talk of?\n",
      "\n",
      "PROSPERO:\n",
      "Obseton pruse, but I knows, noble cross! no, nor heaven,\n",
      "Have I not intello?\n",
      "\n",
      "BUCKINGHAM:\n",
      "Be put thy pawn; am I Cut before us.\n",
      "\n",
      "First Musician:\n",
      "That e'er I shame by us to hell!\n",
      "There, Begggarmen, pent to leadle this\n",
      "a sight hair, anknow you so.\n",
      "\n",
      "TRANIO:\n",
      "That Civerio, both! or I, or thre? then thou drawn!\n",
      "Why.\n",
      "\n",
      "CUTites.\n",
      "\n",
      "C\n"
     ]
    }
   ],
   "source": [
    "def generate_text(model, start_string, num_generate = 1000):\n",
    "    input_eval = [char2idx[ch] for ch in start_string]\n",
    "    input_eval = tf.expand_dims(input_eval, 0)\n",
    "    \n",
    "    text_generated = []\n",
    "    model.reset_states()\n",
    "    \n",
    "    \n",
    "    # 设定一个temperature的值，用于将logits（softmax）之前的数据给扩大或缩小\n",
    "    # temperature > 1, random ，logits的值更加平缓，tf.random.categorical取值更趋向于random\n",
    "    # temperature < 1, greedy ，logits的值更加陡峭，取值更倾向于greedy弹性取值（取最大值）\n",
    "    temperature = 2\n",
    "    \n",
    "    for _ in range(num_generate):\n",
    "        # 1. model inference -> predictions\n",
    "        # 2. sample -> ch -> text_generated.\n",
    "        # 3. update input_eval\n",
    "        \n",
    "        # predictions : [batch_size, input_eval_len, vocab_size]\n",
    "        predictions = model(input_eval)\n",
    "        # predictions: logits -> softmax -> prob\n",
    "        # softmax: e^xi \n",
    "        # eg: 4,2 e^4/(e^4 + e^2) = 0.88, e^2 / (e^4 + e^2) = 0.12\n",
    "        # eg: 2,1 e^2/(e^2 + e) = 0.73, e / (e^2 + e) = 0.27\n",
    "        predictions = predictions / temperature\n",
    "        # predictions : [input_eval_len, vocab_size]\n",
    "        predictions = tf.squeeze(predictions, 0)\n",
    "        # predicted_ids: [input_eval_len, 1]\n",
    "        # a b c -> b c d\n",
    "        predicted_id = tf.random.categorical(\n",
    "            predictions, num_samples = 1)[-1, 0].numpy()\n",
    "        text_generated.append(idx2char[predicted_id])\n",
    "        # s, x -> rnn -> s', y\n",
    "        input_eval = tf.expand_dims([predicted_id], 0)\n",
    "    return start_string + ''.join(text_generated)\n",
    "\n",
    "new_text = generate_text(model2, \"All: \")\n",
    "print(new_text)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
